import numpy as np
import cvxpy as cp


class ConvexOptimization:
    """
    MIMO detection problem: optimize x to minimize norm(H @ x-r)**2 subject to |x_n|=1, arg(x_n) \in A_n, n=1,2,...N
                            H.shape = (M,N), M>=N.
    """

    def __init__(self, N, M, H, r):
        self.N = N  # number of antennas
        self.M = M  # dimensionality of receiving signal
        self.H = H  # wireless channel gain, size = (M,N)
        self.r = r  # receving siganl, size = (M,)
        self.num_solve_problem = 0  # number of solving relaxed optimization problem

    def solve_relaxed_optimization(self, theta):
        """
            @param:
                theta: two-dimensional, the discrete phases of constraint conditions
        """
        assert (len(theta) == self.N)


        x = cp.Variable(self.N, complex=True, name='x')
        obj = cp.Minimize((cp.norm(self.H @ x - self.r, 'fro')))

        constraints = []
        for n in range(self.N):
            if len(theta[n]) == 1:
                constraints += [cp.real(x[n]) == np.cos(theta[n][0])]
                constraints += [cp.imag(x[n]) == np.sin(theta[n][0])]
            elif len(theta[n]) == 2:
                l = 0
                alpha = np.cos((theta[n][l] + theta[n][l + 1]) / 2)
                beta = np.sin((theta[n][l] + theta[n][l + 1]) / 2)
                gamma = np.cos((theta[n][l + 1] - theta[n][l]) / 2)
                constraints += [alpha * cp.real(x[n]) + beta * cp.imag(x[n]) >= gamma]
                constraints += [cp.norm(x[n], 'fro') <= 1]
            else:
                for l in range(len(theta[n]) - 1):
                    alpha = np.cos((theta[n][l] + theta[n][l + 1]) / 2)
                    beta = np.sin((theta[n][l] + theta[n][l + 1]) / 2)
                    gamma = np.cos((theta[n][l + 1] - theta[n][l]) / 2)
                    constraints += [alpha * cp.real(x[n]) + beta * cp.imag(x[n]) <= gamma]
                alpha = np.cos((theta[n][-1] + theta[n][0]) / 2)
                beta = np.sin((theta[n][-1] + theta[n][0]) / 2)
                gamma = np.cos((theta[n][0] - theta[n][-1]) / 2)
                constraints += [alpha * cp.real(x[n]) + beta * cp.imag(x[n]) <= gamma]

        prob = cp.Problem(obj, constraints)
        self.num_solve_problem += 1
        try:
            prob.solve(solver=cp.SCS,verbose=False)
        except Exception as e:
            print(e)
            return np.zeros((self.N,)), np.inf, False

        if prob.status in ['infeasible', 'unbounded']:
            # print('infeasible antenna solution')
            return np.zeros((self.N,)), np.inf, False

        return (x.value, prob.value, True)


if __name__ == '__main__':
    N = 6
    M = N + 1
    H = np.random.randn(M, N) + 1j * np.random.randn(M, N)
    r = np.random.randn(M) + 1j * np.random.randn(M)
    theta = np.random.rand(N, 2) * 2 * np.pi - np.pi
    theta = np.ndarray.tolist(theta)
    # sort theta[n], n=1,2,...N
    for row in range(N):
        theta[row].sort()
    opt = ConvexOptimization(N=N,M=M,H=H,r=r)
    (x, obj, is_solved) = opt.solve_relaxed_optimization(theta=theta)
    print(f"x:{x} \n obj:{obj} \n is_solved:{is_solved}")
